from sentence_transformers import SentenceTransformer

# 封装句子转嵌入向量，返回符合格式的元组，为存入数据库做数据准备

class BegZh:
    def __init__(self,model_path):
        self.embeddings = SentenceTransformer(model_path)

    # 单个句子
    def encode(self,sentence):
        em = self.embeddings.encode(sentence)
        em = [float(f"{x:.18f}") for x in em]
        return em

    # 文档
    def encode_docs(self,sentence,doc_name):
        if isinstance(sentence,str):
            return self.embeddings.encode(sentence)
        elif isinstance(sentence,list):
            result = []
            embeddings_s = self.embeddings.encode(sentence)

            for index,i in enumerate(embeddings_s):
                print(sentence[index])
                i = [float(f"{x:.18f}") for x in i]
                item = (sentence[index],str(i),doc_name)
                result.append(item)

            return result
        else:
            print("入参错误，必须是字符串或字符数组")






if __name__ == '__main__':
    from tools.Pgvector_op import Pgvector
    import configparser

    cf = configparser.ConfigParser()
    cf.read("../config/config.ini")
    beg_zh = BegZh(cf.get("beg_model", "path"))

    print(beg_zh.encode_docs(["你好","哈哈"],'dom.txt'))

    pg = Pgvector(
        cf.get("RAG", "host"),
        cf.get("RAG", "port"),
        cf.get("RAG", "database"),
        cf.get("RAG", "user"),
        cf.get("RAG", "password"),
        cf.get("RAG", "table"),
        mode="RAG"
    )
    pg.creat_table()
    pg.insert(beg_zh.encode_docs(["你好","哈哈"],'dom.txt'))